import argparse
import open3d as o3d
import os
import numpy as np
import tqdm
from vis_utils import get_draw_box

def main(args):
    assert os.path.exists(args.data_path)
    datasets = os.listdir(args.data_path)
    datasets.sort(key=lambda x: x)
    for dataset in tqdm(datasets):
        if "__" == dataset[:2]:#数据集以"__"开头
            dataset_path = os.path.join(args.data_path, dataset)
            main_worker(args, dataset_path)
def main_worker(dataset_path):
    lidar_size = len(os.listdir(dataset_path + '/lidar'))
    radar_size = len(os.listdir(dataset_path + '/radar'))
    label_size = len(os.listdir(dataset_path + '/label_txt'))
    assert (lidar_size == radar_size and lidar_size >= label_size)
    names = []
    for name in os.listdir(dataset_path + '/label_txt'):
        names.append(str.split(name, '.')[0])

    for name in names:

        vis = o3d.visualization.Visualizer()
        vis.create_window(window_name="show_pred_pcd")
        render_option = vis.get_render_option()
        render_option.point_size = 2
        coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=2.0, origin=[0, 0, 0])
        vis.add_geometry(coordinate_frame)

        lidar_path = dataset_path + '/lidar/' + name + ".pcd"
        radar_path = dataset_path + '/radar/' + name + ".pcd"
        label_path = dataset_path + '/label_txt/' + name + ".txt"

        pcd = o3d.io.read_point_cloud(lidar_path)
        lidar = np.asarray(pcd.points)

        datas = o3d.geometry.PointCloud()
        datas.points = o3d.utility.Vector3dVector(lidar[:, :3])
        datas.paint_uniform_color([0.3, 0.5, 0])
        vis.add_geometry(datas)

        pcd = o3d.io.read_point_cloud(radar_path)
        radar = np.asarray(pcd.points)

        data_radar = o3d.geometry.PointCloud()
        data_radar.points = o3d.utility.Vector3dVector(radar[:, :3])
        data_radar.paint_uniform_color([0.3, 0.5, 0])
        vis.add_geometry(data_radar)

        objs = []
        with open(label_path, 'r') as f:
            # open为打开文件，r为读取
            f = open(label_path, 'r')
            # 逐行读取文件内容
            lines = f.readlines()
            for line in lines:
                # 消除空格和空行
                line = line.strip()
                tokens = line.split(' ')
                objs.append(tokens)

        boxes = np.empty(shape=(0, 7))
        for obj in objs:
            box = np.array([[float(obj[1]), float(obj[2]), float(obj[3]), float(obj[4]), float(obj[5]), float(obj[6]),
                             float(obj[7])]])
            boxes = np.concatenate((boxes, box))

        exp_draw_boxes = get_draw_box(boxes)
        for box in exp_draw_boxes:
            vis.add_geometry(box)

        vis.run()
        vis.destroy_window()

    print(0)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--data-path', default="/home/adt/deeplearning/Data/hw0",
                        help='your data root for kitti')

    args = parser.parse_args()

    main(args)
